import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import matplotlib
import matplotlib.dates as mdates
import os

# 设置中文字体（确保系统有中文字体）
matplotlib.rcParams['font.family'] = 'SimHei'
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义保存路径
RESULTS_DIR = os.path.join(os.path.dirname(__file__), '..', 'results')
os.makedirs(RESULTS_DIR, exist_ok=True)  # 确保目录存在


def plot_activity(dau: pd.Series, wau: pd.Series, mau: pd.Series) -> None:
    """
    绘制 DAU / WAU / MAU 活跃用户趋势图，并保存到 results/
    """
    plt.figure(figsize=(14, 6))
    plt.plot(dau.index, dau.values, label="DAU", marker='o', linestyle='-', color='#1f77b4')
    plt.plot(wau.index, wau.values, label="WAU", marker='s', linestyle='--', color='#ff7f0e')
    plt.plot(mau.index, mau.values, label="MAU", marker='^', linestyle='-.', color='#2ca02c')

    plt.title("用户活跃趋势（DAU / WAU / MAU）", fontsize=14)
    plt.xlabel("日期", fontsize=12)
    plt.ylabel("活跃用户数", fontsize=12)
    plt.xticks(rotation=45)
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))

    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.tight_layout()

    # 保存图像
    save_path = os.path.join(RESULTS_DIR, 'activity_trend.png')
    plt.savefig(save_path, dpi=300)
    plt.show()


def plot_retention(retention_df: pd.DataFrame) -> None:
    """
    绘制优化后的用户留存率热力图，并保存到 results/
    """
    df = retention_df.copy()
    df = df.set_index("reg_date").drop(columns=["total_users"], errors='ignore')

    # 调整图像大小
    plt.figure(figsize=(14, 10))

    # 设置最大保留7天，只保留前几列（避免太宽）
    max_days = 7
    df = df.iloc[:, :max_days + 1]

    # 绘制热力图（保留率 * 100, 显示为百分比）
    sns.heatmap(
        df * 100,
        annot=True,
        fmt=".1f",
        cmap="YlGnBu",
        cbar_kws={"label": "留存率 (%)"},
        linewidths=0.5,
        linecolor='gray',
        annot_kws={"fontsize": 8}
    )

    plt.title("用户留存率热力图（前 7 天）", fontsize=16)
    plt.xlabel("留存天数", fontsize=14)
    plt.ylabel("注册日期", fontsize=14)

    # 设置刻度字体大小
    plt.xticks(rotation=0, fontsize=10)
    plt.yticks(rotation=0, fontsize=8)

    plt.tight_layout()

    # 保存图像
    save_path = os.path.join(RESULTS_DIR, 'retention_heatmap.png')
    plt.savefig(save_path, dpi=300)
    plt.show()
